package com.ics.atable.chat.common.milvus;

import com.google.gson.JsonObject;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.*;
import io.milvus.v2.service.vector.request.AnnSearchReq;
import io.milvus.v2.service.vector.request.HybridSearchReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.BaseVector;
import io.milvus.v2.service.vector.request.ranker.RRFRanker;
import io.milvus.v2.service.vector.request.ranker.WeightedRanker;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.SearchResp;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import org.springframework.util.CollectionUtils;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;

import static com.ics.atable.chat.constant.VectorDatabaseConstant.ADDRESS_COLLECTION_NAME;
import static com.ics.atable.chat.constant.VectorDatabaseConstant.ITEM_COLLECTION_NAME;

/**
 * Milvus 多向量场服务
 *
 * @author haohongbin
 */
@Slf4j
@Component
public class MilvusMultiVectorStore {

    @Value("${spring.ai.openai.embedding.options.dimensions}")
    private Integer dimensions;

    private final MilvusClientV2 milvusClientV2;

    public MilvusMultiVectorStore(MilvusClientV2 milvusClientV2) {
        this.milvusClientV2 = milvusClientV2;
    }

    /**
     * 创建 collection
     * @param collectionName collection name
     * @param fieldNameMap field name map
     * @return  Boolean
     */
    public Boolean createCollection(String collectionName, Map<String, List<String>> fieldNameMap) {
        try {
            CreateCollectionReq.CollectionSchema schema = defineMultiVectorSchema(fieldNameMap);
            List<IndexParam> indexParams = defineIndexParams(fieldNameMap.get("dense"), fieldNameMap.get("sparse"));
            CreateCollectionReq createCollectionReq = CreateCollectionReq.builder()
                    .collectionName(collectionName)
                    .collectionSchema(schema)
                    .indexParams(indexParams)
                    .build();
            milvusClientV2.createCollection(createCollectionReq);
            return true;
        } catch (Exception e) {
            log.error("Failed to create collection: {}", collectionName, e);
            return false;
        }
    }

    /**
     * 定义多向量场集合的Schema
     * 包含主键、文本字段、文本密集向量、文本稀疏向量和图像密集向量字段
     * @param fieldNameMap 字段名称映射
     *                     id : 主键
     *                     dense ：稠密向量字段
     *                     sparse : 稀疏向量字段
     *                     metadata : 元数据字段
     */
    private CreateCollectionReq.CollectionSchema defineMultiVectorSchema(Map<String, List<String>> fieldNameMap) {
        // 创建集合Schema构建器
        CreateCollectionReq.CollectionSchema schema = milvusClientV2.createSchema();

        for (Map.Entry<String, List<String>> entry : fieldNameMap.entrySet()){
            switch (entry.getKey()) {
                case "id":
                    schema.addField(AddFieldReq.builder()
                            .fieldName(entry.getValue().getFirst())
                            .dataType(DataType.VarChar)
                            .isPrimaryKey(true)
                            .autoID(false)
                            .maxLength(65535)
                            .build());
                    break;
                case "dense": {
                    for (String fieldName : entry.getValue()) {
                        schema.addField(AddFieldReq.builder()
                                .fieldName(fieldName)
                                .dataType(DataType.FloatVector)
                                .dimension(dimensions)
                                .build());
                    }
                    break;
                }
                case "sparse": {
                    for (String fieldName : entry.getValue()) {
                        schema.addField(AddFieldReq.builder()
                                .fieldName(fieldName)
                                .dataType(DataType.SparseFloatVector)
                                .build());
                    }
                    break;
                }
                case "metadata":
                    schema.addField(AddFieldReq.builder()
                            .fieldName(entry.getValue().getFirst())
                            .dataType(DataType.JSON)
                            .build());
                    break;
            }
        }
        return schema;
    }


    /**
     * 创建索引
     * @param sparseFieldNames 稀疏字段名称
     * @param denseFieldNames 稠密字段名称
     * @return IndexParam List
     */
    private List<IndexParam> defineIndexParams(List<String> sparseFieldNames, List<String> denseFieldNames) {

        List<IndexParam> indexParams = new ArrayList<>();

        for (String fieldName : denseFieldNames) {
            IndexParam indexParamForTextDense = IndexParam.builder()
                    .fieldName(fieldName)
                    .indexName(fieldName + "_dense_index")
                    .indexType(IndexParam.IndexType.AUTOINDEX)
                    .metricType(IndexParam.MetricType.COSINE)
                    .build();
            indexParams.add(indexParamForTextDense);
        }

        for (String fieldName : sparseFieldNames) {
            Map<String, Object> sparseParams = new HashMap<>();
            sparseParams.put("inverted_index_algo", "DAAT_MAXSCORE");
            IndexParam indexParamForTextSparse = IndexParam.builder()
                    .fieldName(fieldName)
                    .indexName(fieldName + "_sparse_index")
                    .indexType(IndexParam.IndexType.SPARSE_INVERTED_INDEX)
                    .metricType(IndexParam.MetricType.BM25)
                    .extraParams(sparseParams)
                    .build();
            indexParams.add(indexParamForTextSparse);
        }

        return indexParams;
    }



    /**
     * 检查集合是否存在
     * @param collectionName 集合名称
     * @return boolean 集合是否存在
     */
    public boolean hasCollection(String collectionName) {
        try {
            HasCollectionReq req = HasCollectionReq.builder()
                    .collectionName(collectionName)
                    .build();
            return Boolean.TRUE.equals(milvusClientV2.hasCollection(req));
        } catch (Exception e) {
            log.error("检查集合存在性时出错: {}", collectionName, e);
            return false;
        }
    }

    /**
     * 删除集合
     * @param collectionName 集合名称
     * @return boolean 删除是否成功
     */
    public boolean dropCollection(String collectionName) {
        try {
            if (hasCollection(collectionName)) {
                DropCollectionReq req = DropCollectionReq.builder()
                        .collectionName(collectionName)
                        .build();
                milvusClientV2.dropCollection(req);
                log.info("成功删除集合: {}", collectionName);
                return true;
            } else {
                log.warn("集合不存在: {}", collectionName);
                return false;
            }
        } catch (Exception e) {
            log.error("删除集合时出错: {}", collectionName, e);
            return false;
        }
    }

    /**
     * 加载集合到内存
     * @param collectionName 集合名称
     */
    public void loadCollection(String collectionName) {
        try {
            LoadCollectionReq loadCollectionReq = LoadCollectionReq.builder()
                    .collectionName(collectionName).
                    build();
            milvusClientV2.loadCollection(loadCollectionReq);
            log.info("成功加载集合到内存: {}", collectionName);
        } catch (Exception e) {
            log.error("加载集合到内存时出错: {}", collectionName, e);
        }
    }

    /**
     * 从内存中释放集合
     * @param collectionName 集合名称
     */
    public void releaseCollection(String collectionName) {
        try {
            ReleaseCollectionReq releaseCollectionReq = ReleaseCollectionReq.builder()
                    .collectionName(collectionName)
                    .build();
            milvusClientV2.releaseCollection(releaseCollectionReq);
            log.info("成功从内存中释放集合: {}", collectionName);
        } catch (Exception e) {
            log.error("从内存中释放集合时出错: {}", collectionName, e);
        }
    }

    /**
     * 插入数据
     * @param data 数据列表
     * @param collectionName 集合名称
     * @return InsertResp 插入响应
     */
    public InsertResp insert(List<JsonObject> data, String collectionName) {
        try {
            InsertReq insertReq = InsertReq.builder()
                    .collectionName(collectionName)
                    .data(data)
                    .build();
            InsertResp insertResp = milvusClientV2.insert(insertReq);
            log.info("成功插入 {} 条数据到集合: {}", data.size(), collectionName);
            return insertResp;
        } catch (Exception e) {
            log.error("插入数据时出错，集合: {}", collectionName, e);
            throw new RuntimeException("插入数据失败", e);
        }
    }

    /**
     * 异步插入数据
     * @param data 数据列表
     * @param collectionName 集合名称
     * @param batchSize 批次大小
     * @return CompletableFuture<List<InsertResp>> 插入响应列表
     */
    public CompletableFuture<List<InsertResp>> insertAsync(List<JsonObject> data, String collectionName, int batchSize) {
        ExecutorService executor = Executors.newFixedThreadPool(5); // 可根据需要调整线程数

        List<CompletableFuture<InsertResp>> futures = new ArrayList<>();

        // 将数据分批
        for (int i = 0; i < data.size(); i += batchSize) {
            int endIndex = Math.min(i + batchSize, data.size());
            List<JsonObject> batch = data.subList(i, endIndex);

            CompletableFuture<InsertResp> future = CompletableFuture.supplyAsync(() -> {
                try {
                    InsertReq insertReq = InsertReq.builder()
                            .collectionName(collectionName)
                            .data(batch)
                            .build();
                    InsertResp insertResp = milvusClientV2.insert(insertReq);
                    log.info("成功异步插入 {} 条数据到集合: {}", batch.size(), collectionName);
                    return insertResp;
                } catch (Exception e) {
                    log.error("异步插入数据时出错，集合: {}", collectionName, e);
                    throw new RuntimeException("异步插入数据失败", e);
                }
            }, executor);

            futures.add(future);
        }

        // 等待所有插入操作完成
        return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
                .thenApply(v -> futures.stream()
                        .map(CompletableFuture::join)
                        .collect(Collectors.toList()))
                .whenComplete((result, throwable) -> executor.shutdown());
    }


    /**
     * 混合搜索
     * @param collectionName 集合名称
     * @param queryVectors 查询向量列表
     * @param denseFieldNames 稠密向量字段名列表
     * @param weights 字段权重列表
     * @param topK 返回结果数
     * @param filter 过滤条件表达式
     * @return SearchResp 搜索响应
     */
    public SearchResp hybridSearch(String collectionName, List<BaseVector> queryVectors, List<String> denseFieldNames,
                                   List<Float> weights, int topK, String filter) {
        try {
            List<AnnSearchReq> searchRequests = new ArrayList<>();
            for (String fieldName : denseFieldNames) {
               AnnSearchReq.AnnSearchReqBuilder req = AnnSearchReq.builder()
                        .vectorFieldName(fieldName)
                        .vectors(queryVectors)
                        .params("{\"nprobe\": 10}")
                        .topK(topK);
               if (StringUtils.isNotBlank(filter)) {
                   req.filter(filter);
               }
                searchRequests.add(req.build());
            }

            // 根据是否提供权重选择重排器
            CreateCollectionReq.Function ranker;
            if (!CollectionUtils.isEmpty(weights)) {
                if (denseFieldNames.size() != weights.size()) {
                    throw new IllegalArgumentException("字段名称数量与权重数量不匹配");
                }
                ranker = new WeightedRanker(weights);
            } else {
                ranker = new RRFRanker(100);
            }


            HybridSearchReq hybridSearchReq = HybridSearchReq.builder()
                    .collectionName(collectionName)
                    .searchRequests(searchRequests)
                    .ranker(ranker)
                    .topK(topK)
                    .outFields(Arrays.asList("*"))
                    .build();

            SearchResp searchResp = milvusClientV2.hybridSearch(hybridSearchReq);
            String rankerType = (ranker instanceof WeightedRanker) ? "加权重排器" : "RRF重排器";
            log.info("混合搜索完成（使用{}），集合: {}, 返回结果数: {}",
                    rankerType, collectionName, searchResp.getSearchResults().size());
            return searchResp;
        } catch (Exception e) {
            log.error("混合搜索时出错，集合: {}", collectionName, e);
            throw new RuntimeException("搜索失败", e);
        }
    }

    /**
     * 向量搜索
     * @param collectionName 集合名称
     * @param queryVectors 查询向量列表
     * @param denseFieldNames 稠密向量字段名列表
     * @param topK 返回结果数
     * @param filter 过滤条件表达式
     * @return SearchResp 搜索响应
     */
    public SearchResp aVectorSearch(String collectionName, List<BaseVector> queryVectors, String denseFieldNames,
                                     int topK, String filter) {
        try {
            SearchReq searchReq = SearchReq.builder()
                    .databaseName("atable")
                    .collectionName(collectionName)
                    .data(queryVectors)
                    .topK(topK)
                    .filter(filter)
                    .outputFields(Collections.singletonList("*"))
                    .build();
            return milvusClientV2.search(searchReq);
        } catch (Exception e) {
            log.error("向量搜索时出错，集合: {}", collectionName, e);
            throw new RuntimeException("搜索失败", e);
        }
    }


    /**
     * 获取集合统计信息
     * @param collectionName 集合名称
     * @return long 集合中的实体数量
     */
    public long getCollectionEntityCount(String collectionName) {
        try {
            GetCollectionStatsReq getCollectionStatsReq = GetCollectionStatsReq.builder()
                    .collectionName(collectionName)
                    .build();
            return milvusClientV2.getCollectionStats(getCollectionStatsReq).getNumOfEntities();
        } catch (Exception e) {
            log.error("获取集合统计信息时出错: {}", collectionName, e);
            return -1;
        }
    }


    public void preloadCollections() {
        try {
            loadCollection(ITEM_COLLECTION_NAME);
            loadCollection(ADDRESS_COLLECTION_NAME);
            log.info("预加载 Milvus collections 完成: {}, {}", ITEM_COLLECTION_NAME, ADDRESS_COLLECTION_NAME);
        } catch (Exception e) {
            log.error("预加载 Milvus collections 失败", e);
            throw new RuntimeException("Milvus collections 预加载失败", e);
        }
    }

    /**
     * 初始化方法，在Bean初始化后调用
     */
    @PostConstruct
    public void initialize() {
        preloadCollections();
    }

}
